{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Carlane: Experiment Code\n",
    "\n",
    "train and test mask rcnn model of Carlane dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import sys\n",
    "ROOT_DIR = os.path.abspath(\"../code\")\n",
    "sys.path.append(ROOT_DIR)\n",
    "import carlane\n",
    "\n",
    "from mrcnn import model as modellib, utils\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "logs = '../log'\n",
    "if not os.path.exists(logs):\n",
    "    os.mkdir(logs)\n",
    "    \n",
    "command = \"train\" #  or val\n",
    "dataset_dir = \"../data/2.7万张道路线精细标注抽取2000张\"\n",
    "\n",
    "weights = \"coco\" #coco, or last, or imagenet\n",
    "COCO_WEIGHTS_PATH = '../saved_model/mask_rcnn_coco.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if command == \"train\":\n",
    "    config = carlane.CarlaneConfig()\n",
    "else:\n",
    "    class InferenceConfig(CocoConfig):\n",
    "        # Set batch size to 1 since we'll be running inference on\n",
    "        # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU\n",
    "        GPU_COUNT = 1\n",
    "        IMAGES_PER_GPU = 1\n",
    "        DETECTION_MIN_CONFIDENCE = 0\n",
    "    config = InferenceConfig()\n",
    "config.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if command == \"train\":\n",
    "    model = modellib.MaskRCNN(mode=\"training\", config=config,\n",
    "                              model_dir=logs)\n",
    "else:\n",
    "    model = modellib.MaskRCNN(mode=\"inference\", config=config,\n",
    "                              model_dir=logs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if weights.lower() == \"imagenet\":\n",
    "    # Start from ImageNet trained weights\n",
    "    weights_path = model.get_imagenet_weights()\n",
    "elif weights.lower() == \"coco\":\n",
    "    weights_path = COCO_WEIGHTS_PATH\n",
    "else:\n",
    "    weights_path = weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Loading weights \", weights_path)\n",
    "if weights.lower() == \"coco\":\n",
    "    # Exclude the last layers because they require a matching\n",
    "    # number of classes\n",
    "    model.load_weights(weights_path, by_name=True, exclude=[\n",
    "        \"mrcnn_class_logits\", \"mrcnn_bbox_fc\",\n",
    "        \"mrcnn_bbox\", \"mrcnn_mask\"])\n",
    "else:\n",
    "    model.load_weights(weights_path, by_name=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Training dataset.\n",
    "dataset_train = carlane.CarlaneDataset()\n",
    "dataset_train.load_carlane(dataset_dir, \"train\")\n",
    "dataset_train.prepare()\n",
    "\n",
    "# Validation dataset\n",
    "dataset_val = carlane.CarlaneDataset()\n",
    "dataset_val.load_carlane(dataset_dir, \"val\")\n",
    "dataset_val.prepare()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"train images count:{}, val images count: {}\".format(dataset_train.num_images, dataset_val.num_images))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import imgaug\n",
    "# Image Augmentation\n",
    "# Right/Left flip 50% of the time\n",
    "augmentation = imgaug.augmenters.Fliplr(0.5)\n",
    "\n",
    "# *** This training schedule is an example. Update to your needs ***\n",
    "\n",
    "if weights.lower() == 'imagenet':\n",
    "    # Training - Stage 1\n",
    "    print(\"Training network heads\")\n",
    "    model.train(dataset_train, dataset_val,\n",
    "                learning_rate=config.LEARNING_RATE,\n",
    "                epochs=40,\n",
    "                layers='heads',\n",
    "                augmentation=augmentation)\n",
    "\n",
    "'''\n",
    "# Training - Stage 2\n",
    "# Finetune layers from ResNet stage 4 and up\n",
    "print(\"Fine tune Resnet stage 4 and up\")\n",
    "model.train(dataset_train, dataset_val,\n",
    "            learning_rate=config.LEARNING_RATE,\n",
    "            epochs=80,\n",
    "            layers='4+',\n",
    "            augmentation=augmentation)\n",
    "'''\n",
    "# Training - Stage 3\n",
    "# Fine tune all layers\n",
    "print(\"Fine tune all layers\")\n",
    "model.train(dataset_train, dataset_val,\n",
    "            learning_rate=config.LEARNING_RATE,\n",
    "            epochs=120,\n",
    "            layers='all',\n",
    "            augmentation=augmentation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Draft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "# Creates a graph.\n",
    "a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')\n",
    "b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')\n",
    "c = tf.matmul(a, b)\n",
    "# Creates a session with log_device_placement set to True.\n",
    "sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))\n",
    "# Runs the op.\n",
    "print(sess.run(c))\n",
    "\n",
    "# Creates a graph.\n",
    "a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')\n",
    "b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')\n",
    "c = tf.matmul(a, b)\n",
    "# Creates a session with log_device_placement set to True.\n",
    "sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))\n",
    "# Runs the op.\n",
    "print(sess.run(c))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
